use rust_tcp::{Quad, IPV4_PROTOCOL, TCB, TCP_PROTOCOL};
use std::{collections::HashMap, io::Result};
use tun_tap::{Iface, Mode};

fn main() -> Result<()> {
    let mut iface = Iface::new("mytun", Mode::Tun).expect("Failed to create a TUN device");
    // let _name = iface.name();
    // Configure the device ‒ set IP address on it, bring it up.
    let mut buffer = vec![0u8; 1504]; // MTU + 4 for the header
    let mut connections: HashMap<Quad, TCB> = Default::default();
    loop {
        let nbytes = iface.recv(&mut buffer).unwrap();
        let eth_proto = u16::from_be_bytes([buffer[2], buffer[3]]);
        if eth_proto != IPV4_PROTOCOL {
            // 如果不是IPV4协议，不予处理
            continue;
        }
        match etherparse::Ipv4HeaderSlice::from_slice(&buffer[4..nbytes]) {
            Ok(ip_header) => {
                if ip_header.protocol() != TCP_PROTOCOL {
                    // 如果不是TCP协议，不予处理
                    continue;
                }
                let src_ip = ip_header.source_addr();
                let dst_ip = ip_header.destination_addr();
                match etherparse::TcpHeaderSlice::from_slice(
                    &buffer[4 + ip_header.slice().len()..nbytes],
                ) {
                    Ok(tcp_header) => {
                        use std::collections::hash_map::Entry;
                        let src_port = tcp_header.source_port();
                        let dst_port = tcp_header.destination_port();
                        let data_idx = 4 + ip_header.slice().len() + tcp_header.slice().len();
                        // 判断map中是否已有该连接，如果没有，则建立，否则正常通信
                        match connections.entry(Quad {
                            src: (src_ip, src_port),
                            dst: (dst_ip, dst_port),
                        }) {
                            Entry::Occupied(mut tcb) => {
                                tcb.get_mut().receive_packet(&mut iface, ip_header, tcp_header, &buffer[data_idx..nbytes])?;
                            }
                            Entry::Vacant(entry) => {
                                if let Some(tcb) = TCB::accept(
                                    &mut iface,
                                    ip_header,
                                    tcp_header,
                                    &buffer[data_idx..nbytes]
                                )? {
                                    entry.insert(tcb);
                                }
                            }
                        }
                    }
                    Err(e) => {
                        eprintln!("Failed to parse TCP header: {}", e);
                        continue;
                    }
                }
            }
            Err(e) => {
                eprintln!("Failed to parse IPv4 header: {}", e);
                continue;
            }
        }
    }
    // Ok(())
}
